import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt

df = pd.read_excel('edges_output.xlsx')

# 过滤掉 f = 0 的数据
df_hehe = df[df['f'] != 0]

u = df_hehe['u']
v = df_hehe['v']
f = df_hehe['f']

G = nx.DiGraph()

# 使用 .iloc 通过位置访问数据
for i in range(len(u)):  # 遍历所有数据行
    if u.iloc[i] == v.iloc[i]:  # 判断是否是自环
        print('Self-loop detected at:', u.iloc[i])
    else:
        if not G.has_node(u.iloc[i]):
            G.add_node(u.iloc[i])
        if not G.has_node(v.iloc[i]):
            G.add_node(v.iloc[i])
        G.add_edge(u.iloc[i], v.iloc[i], capacity=f.iloc[i])

# 获取边的颜色列表
edge_color_list = [(u, v, d['capacity']) for u, v, d in G.edges(data=True)]

# 计算容量值的最大值
capacities = [capacity for u, v, capacity in edge_color_list]
max_capacity = max(capacities)

# 设置边的颜色
edge_colors = [plt.cm.autumn(capacity / max_capacity) for u, v, capacity in edge_color_list]

# 绘制图形
nx.draw(G, edge_color=edge_colors)
plt.show()

